import asyncio
import os
import signal
from typing import Optional
import typing as tp

from loguru import logger


class AIOCmdExecutor:
    CMD_ERROR_CODE = -9999
    TIMEOUT_CODE = -9998
    INTERRUPT_CODE = -9997
    UNKNOWN_CODE = -9996

    def __init__(
            self,
            cmd: list,
            timeout: Optional[float] = None,
            as_shell: bool = False,
            log_callback: tp.Callable[[str, tp.Literal["STDERR", "STDOUT"]], tp.Coroutine] = None,
            **kwargs
    ):
        self.cmd = cmd
        self.timeout = timeout
        self.as_shell = as_shell
        self.log_cb = log_callback
        self.kwargs = kwargs
        self.p: Optional[asyncio.subprocess.Process] = None
        self.out_task: Optional[asyncio.Task] = None
        self.err_task: Optional[asyncio.Task] = None

        self.is_stop = False
        self.is_interrupt = False

    async def __aenter__(self) -> "AIOCmdExecutor":
        try:
            if self.as_shell:
                self.p = await asyncio.create_subprocess_shell(
                    " ".join(self.cmd),
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                    preexec_fn=os.setpgrp,
                    **self.kwargs
                )
            else:
                self.p = await asyncio.create_subprocess_exec(
                    *self.cmd,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                    preexec_fn=os.setpgrp,
                    **self.kwargs
                )
        except (BaseException, Exception):
            import traceback
            logger.error(traceback.format_exc())
        else:
            self.out_task = asyncio.create_task(self._read_stream(self.p.stdout, "STDOUT"))
            self.err_task = asyncio.create_task(self._read_stream(self.p.stderr, "STDERR"))
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        self.out_task = None
        self.err_task = None
        self.p = None

    async def run(self) -> int:
        if self.p is None:
            return self.CMD_ERROR_CODE

        try:
            await asyncio.wait_for(self.p.wait(), timeout=self.timeout)
            code = self.p.returncode
        except asyncio.TimeoutError:
            self.p.kill()
            code = self.TIMEOUT_CODE
        except (BaseException, Exception):
            import traceback
            logger.error(traceback.format_exc)
            code = self.UNKNOWN_CODE

        if self.is_interrupt:
            code = self.INTERRUPT_CODE
        self.is_stop = True
        return code

    async def _read_stream(self, stream: asyncio.StreamReader, type_: tp.Literal["STDERR", "STDOUT"]):
        while True:
            if self.is_stop:
                break

            line = await stream.readline()
            if not line:
                break
            decoded = line.decode('utf-8', errors='ignore').strip()
            if self.log_cb:
                await self.log_cb(decoded, type_)
            else:
                print(decoded)

    async def interrupt(self):
        if self.p:
            self.is_interrupt = True
            pgid = os.getpgid(self.p.pid)
            os.killpg(pgid, signal.SIGKILL)
            if self.p and self.log_cb:
                await self.log_cb(F"pgid: {pgid}, pid: {self.p.pid} 已终止", "STDOUT")


class CMDExecutor(AIOCmdExecutor):
    def __init__(
            self,
            cmd: list,
            log_callback: tp.Callable[[str, tp.Literal["STDERR", "STDOUT"]], tp.Coroutine] = None,
            timeout: Optional[float] = None,
            **kwargs
    ):
        super().__init__(cmd, timeout, as_shell=False, log_callback=log_callback, **kwargs)


class ShellExecutor(AIOCmdExecutor):
    def __init__(
            self,
            cmd: list,
            log_callback: tp.Callable[[str, tp.Literal["STDERR", "STDOUT"]], tp.Coroutine] = None,
            timeout: Optional[float] = None,
            **kwargs
    ):
        super().__init__(cmd, timeout, as_shell=True, log_callback=log_callback, **kwargs)


if __name__ == "__main__":
    async def main():
        async with CMDExecutor([r"D:\my_py\mnagent\.venv\Scripts\python.exe", "-c", "while True: pass"], timeout=5) as exe:

            async def interrupt():
                await asyncio.sleep(2)
                exe.interrupt()

            asyncio.create_task(interrupt())
            print(await exe.run())

    asyncio.run(main())
